#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import rospy
import rospkg
import yaml
import subprocess
import signal
import os
from service_launcher.srv import LaunchControl, LaunchControlResponse

class LaunchManager:
    def __init__(self):
        rospy.init_node('launch_manager', anonymous=True)

        # --- Parameters ---
        config_path_param = rospy.get_param('~config_path', 'config/launches.yaml')
        
        # --- Load launch configurations ---
        self.rospack = rospkg.RosPack()
        self.config_path = self._resolve_package_path(config_path_param)
        if not self.config_path:
            rospy.logfatal("Could not find config file at: %s", config_path_param)
            return

        self.launch_configs = self._load_config(self.config_path)
        if self.launch_configs is None:
            rospy.logfatal("Failed to load or parse config file.")
            return
            
        rospy.loginfo("Launch Manager initialized. Loaded %d launch configurations.", len(self.launch_configs))

        # --- State Management ---
        self.active_launches = {}  # {launch_name: subprocess.Popen object}

        # --- ROS Interface ---
        self.service = rospy.Service('/manage_launch', LaunchControl, self.handle_launch_control)
        rospy.on_shutdown(self.shutdown)

        rospy.loginfo("Launch Manager is ready to receive requests.")

    def _resolve_package_path(self, path):
        """ Resolves a path that might be relative to a package. """
        if not path.startswith('/'):
            try:
                # Assuming the path is relative to the package root, e.g., "config/launches.yaml"
                # This needs a package context, let's assume it's in this package
                pkg_path = self.rospack.get_path('service_launcher')
                return os.path.join(pkg_path, path)
            except rospkg.ResourceNotFound:
                rospy.logerr("Could not find package 'service_launcher' to resolve path: %s", path)
                return None
        return path

    def _load_config(self, path):
        """ Loads the YAML configuration file. """
        try:
            with open(path, 'r') as f:
                config_data = yaml.safe_load(f)
                # Create a mapping from name to config for quick access
                return {item['name']: item for item in config_data.get('launches', [])}
        except Exception as e:
            rospy.logerr("Error loading config file %s: %s", path, e)
            return None

    def handle_launch_control(self, req):
        """ Handles start/stop requests from the ROS service. """
        launch_name = req.launch_name
        command = req.command.lower()

        if launch_name not in self.launch_configs:
            rospy.logwarn("Request for unknown launch configuration: %s", launch_name)
            return LaunchControlResponse(success=False, message="Launch configuration not found.", status="not_found")

        if command == "start":
            return self._start_launch(launch_name)
        elif command == "stop":
            return self._stop_launch(launch_name)
        else:
            rospy.logwarn("Invalid command received: %s", command)
            return LaunchControlResponse(success=False, message="Invalid command. Use 'start' or 'stop'.", status="error")

    def _start_launch(self, launch_name):
        """ Starts a roslaunch process. """
        if launch_name in self.active_launches:
            proc = self.active_launches[launch_name]
            if proc.poll() is None:
                rospy.logwarn("Launch '%s' is already running with PID %d.", launch_name, proc.pid)
                return LaunchControlResponse(success=False, message="Already running.", pid=proc.pid, status="running")
        
        config = self.launch_configs[launch_name]
        package = config['package']
        launch_file = config['launch_file']
        
        command = ["roslaunch", package, launch_file]
        
        try:
            rospy.loginfo("Starting launch: %s", " ".join(command))
            # Use preexec_fn=os.setsid to create a new process group.
            # This allows us to kill the entire process tree easily.
            proc = subprocess.Popen(command, preexec_fn=os.setsid)
            self.active_launches[launch_name] = proc
            rospy.loginfo("Launch '%s' started successfully with PID %d.", launch_name, proc.pid)
            return LaunchControlResponse(success=True, message="Process started.", pid=proc.pid, status="running")
        except Exception as e:
            rospy.logerr("Failed to start launch '%s': %s", launch_name, e)
            return LaunchControlResponse(success=False, message=str(e), status="error")

    def _stop_launch(self, launch_name):
        """ Stops a roslaunch process. """
        if launch_name not in self.active_launches:
            rospy.logwarn("Request to stop non-running launch: %s", launch_name)
            return LaunchControlResponse(success=False, message="Not running or not managed by this manager.", status="stopped")

        proc = self.active_launches.pop(launch_name)
        pid = proc.pid
        
        if proc.poll() is not None:
             rospy.logwarn("Process for '%s' (PID %d) was already terminated.", launch_name, pid)
             return LaunchControlResponse(success=True, message="Process was already terminated.", pid=pid, status="stopped")

        try:
            rospy.loginfo("Stopping launch '%s' (PID %s)...", launch_name, pid)
            # Send SIGINT to the entire process group. This allows roslaunch to gracefully shut down its children.
            os.killpg(os.getpgid(pid), signal.SIGINT)
            proc.wait(timeout=10) # Wait for graceful shutdown
            rospy.loginfo("Launch '%s' (PID %s) stopped gracefully.", launch_name, pid)
            return LaunchControlResponse(success=True, message="Process stopped.", pid=pid, status="stopped")
        except subprocess.TimeoutExpired:
            rospy.logwarn("Graceful shutdown for '%s' (PID %s) timed out. Forcing kill...", launch_name, pid)
            os.killpg(os.getpgid(pid), signal.SIGKILL)
            rospy.logwarn("Launch '%s' (PID %s) force killed.", launch_name, pid)
            return LaunchControlResponse(success=True, message="Process force killed after timeout.", pid=pid, status="stopped")
        except Exception as e:
            rospy.logerr("Error stopping launch '%s' (PID %s): %s", launch_name, pid, e)
            return LaunchControlResponse(success=False, message=str(e), pid=pid, status="error")

    def shutdown(self):
        """ Cleans up all running processes on node shutdown. """
        rospy.loginfo("Launch Manager shutting down. Stopping all active launches...")
        # Create a copy of keys because the dictionary will be modified during iteration
        active_launch_names = list(self.active_launches.keys())
        for name in active_launch_names:
            self._stop_launch(name)
        rospy.loginfo("All launches stopped. Shutdown complete.")

if __name__ == '__main__':
    try:
        LaunchManager()
        rospy.spin()
    except rospy.ROSInterruptException:
        pass